package com.example.userservice.client;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.ResultCallback;

import java.util.concurrent.CompletableFuture;

public class DashScopeClient {

    private final Generation generation;

    public DashScopeClient() {
        generation = new Generation();
    }

    /**
     * 同步调用ai模型
     */
    public String callSync(String prompt, String model) {
        try {
            GenerationParam param = GenerationParam.builder()
                    .model(model)
                    .prompt(prompt)
                    .build();
            GenerationResult result = generation.call(param);
//            return result.getOutput().getChoices().get(0).getMessage().getContent();
//            return JsonUtils.toJson(result);
            return result.getOutput().getText();
        } catch (Exception e) {
            throw new RuntimeException("AI调用失败", e);
        }
    }

    /**
     * 异步调用ai模型
     */
    public CompletableFuture<String> callAsync(String prompt, String model) {
        CompletableFuture<String> future = new CompletableFuture<>();
        try {
            GenerationParam param = GenerationParam.builder()
                    .model(model)
                    .prompt(prompt)
                    .build();

            generation.call(param, new ResultCallback<GenerationResult>() {
                @Override
                public void onEvent(GenerationResult result) {
                    String content = result.getOutput().getText();
                    future.complete(content);
                }

                @Override
                public void onComplete() {

                }

                @Override
                public void onError(Exception e) {
                    future.completeExceptionally(e);
                }
            });
        } catch (Exception e) {
            future.completeExceptionally(e);
        }

        return future;
    }

    /**
     * 带参数的高级调用
     */
    public String callWithParam(String prompt, String model,
                                Double temperature, Integer maxTokens) {
        try {
            GenerationParam.GenerationParamBuilder builder = GenerationParam.builder()
                    .model(model)
                    .prompt(prompt);

            if (temperature != null) {
                builder.temperature(Float.valueOf(String.valueOf(temperature)));
            }
            if (maxTokens != null) {
                builder.maxTokens(maxTokens);
            }

            GenerationResult result = generation.call(builder.build());
            return result.getOutput().getText();
        } catch (Exception e) {
            throw new RuntimeException("AI调用失败", e);
        }
    }
}
